from export.export import export_to_excel
from rapiddns.rapiddns import RapidDnsScan
from thirdapi.crtsh import CrtSh
from thirdapi.virustotal import VirustotalScan
from dnsx.dnsx import DNSXScan
from config.settings import PACKAGE_PATH
from lib.regex import is_domain
import asyncio
import pandas as pd


class SubdomainScan:
    def __init__(self, domain):
        self.domain = domain
        self.rapid = RapidDnsScan(domain)
        self.crtsh = CrtSh()
        self.virustotal = VirustotalScan()
        self.results = []
        self.crt_results = []
    
    def rapid_scan(self):
        self.rapid.scan()
        self.results.extend(self.rapid.results)
        
    
    def third_api_scan(self):
        dnsx_exec_path = PACKAGE_PATH / 'thirdparty/dnsx'  
        self.crtsh.scan(self.domain)
        self.crt_results.extend(self.crtsh.raw)
        self.virustotal.scan(self.domain)
        results = list(self.crtsh.results.union(self.virustotal.results))
        dnsx = DNSXScan(dnsx_exec_path, results, query=self.domain)
        asyncio.run(dnsx.scan())
        self.results.extend(dnsx.results)
        
    def scan(self):
        self.rapid_scan()
        self.third_api_scan()
    
    def export(self, save_path=None):
        results = pd.DataFrame.from_dict(self.results)
        results = results.drop_duplicates(subset=['subdomain', 'address'], keep='first')
        columns_dict = {
            'subdomain': '子域名',
            'address': '解析结果',
            'type': '解析类型',
            'test': '测试'
        }
        export_to_excel(results, xlsx_name=self.domain, sheet_name='domain-' +
                        self.domain, columns_dict=columns_dict, save_path=save_path)


class Subdomain:
    def __init__(self, domains):
        self.domains = domains
        self.crt_results = []
    
    def run(self, save_path):
        for domain in self.domains:
            if not is_domain(domain):
                continue
            subdomain = SubdomainScan(domain)
            subdomain.scan()
            self.crt_results.extend(subdomain.crt_results)
            print(f'解析子域名数量: {len(subdomain.results)}')
            if len(subdomain.results) == 0:
                continue
            subdomain.export(save_path)
        export_to_excel(self.crt_results, xlsx_name='', sheet_name='证书信息', save_path=save_path) 